import qualified Data.Set as Set

type Jolt = Int
data Adapters = Adapters (Set.Set Jolt) Jolt
type State = (Jolt, (Int,Int,Int))

parseInput :: String -> Adapters
parseInput string = 
    Adapters (Set.fromList $ (max+3):numbers) (max+3)
    where
        numbers = map read . lines $ string
        max = maximum numbers

step :: Adapters -> State -> State
step adapters@(Adapters numbers max) state@(jolt, (size1,size2,size3)) = 
    if jolt == max
        then state
        else step adapters newState
    where
        newState 
            | Set.member (jolt+1) numbers = (jolt+1, (size1+1,size2,size3))
            | Set.member (jolt+2) numbers = (jolt+2, (size1,size2+1,size3))
            | Set.member (jolt+3) numbers = (jolt+3, (size1,size2,size3+1))
            | otherwise = error "No jolt"

getJoltDifferences :: Adapters -> (Int,Int,Int)
getJoltDifferences adapters@(Adapters numbers max) = sizes
    where
        (_,sizes) = step adapters (0, (0,0,0))

getProduct :: (Int,Int,Int) -> Int
getProduct (size1,_,size3) = size1 * size3

main :: IO ()
main = do
    input <- readFile "input.txt"
    print . getProduct . getJoltDifferences . parseInput $ input
